use crate::{
    transport::{McpMessage, McpTransport}, McpClientInitializeResult, McpError, McpErrorCode, McpRequest, McpResponse, McpResponseError, McpResult, McpServerInitializeResult, McpTool, McpToolCall, McpToolCallResult
};
use async_trait::async_trait;
use futures::StreamExt;
use log::debug;
use std::sync::Arc;
use tokio::sync::RwLock;

#[async_trait]
pub trait McpServerHandler: Send + Sync {
    async fn initialize(&self, result: McpClientInitializeResult) -> McpResult<McpServerInitializeResult>;

    async fn tools_list(&self) -> McpResult<Vec<McpTool>>;
    async fn tools_call(&self, params: Option<serde_json::Value>) -> McpResult<McpToolCallResult> {
        let params = params.ok_or_else(|| {
            McpError::protocol(McpErrorCode::InvalidParams, "Missing parameters in tools/call request".to_string())
        })?;

        // Parse the parameters as CallToolParams
        let call_params: McpToolCall = serde_json::from_value(params.clone())
            .map_err(|e| McpError::protocol(McpErrorCode::InvalidParams, format!("Invalid tools/call parameters: {}", e)))?;
        
        let tool_name = &call_params.name;
        self.execute_tool(&tool_name, call_params.arguments).await
    }

    async fn execute_tool(
        &self,
        method: &str,
        arguments: Option<serde_json::Value>,
    ) -> McpResult<McpToolCallResult>;

    async fn shutdown(&self) -> McpResult<()>;
}

pub struct McpServer {
    transport: Arc<dyn McpTransport>,
    handler: Arc<dyn McpServerHandler>,
    initilized: Arc<RwLock<bool>>,
}

impl McpServer {
    pub fn new(transport: Arc<dyn McpTransport>, handler: Arc<dyn McpServerHandler>) -> Self {
        Self {
            transport,
            handler,
            initilized: Arc::new(RwLock::new(false))
        }
    }

    pub async fn run(&self) -> McpResult<()> {
        let mut stream = self.transport.receive();

        while let Some(message) = stream.next().await {
            match message? {
                McpMessage::Request(request) => {
                    let response = match self.handle_request(request.clone()).await {
                        Ok(response) => response,
                        Err(err) => McpResponse::error(request.id, McpResponseError::from(err)),
                    };
                    
                    self.transport.send(McpMessage::Response(response)).await?;
                }
                McpMessage::Notification(_notification) => {}
                McpMessage::Response(_) => {
                    return Err(McpError::protocol(
                        McpErrorCode::InvalidRequest,
                        "Server received unepected response",
                    ))
                }
            }
        }

        Ok(())
    }

    pub async fn handle_request(&self, request: McpRequest) -> McpResult<McpResponse> {    
        debug!("Handle a request '{}' with '{:?}'", request.method, request.params);
        match request.method.as_str() {
            "initialize" => {
                self.check_initialized(false).await?;

                let params = request.params.ok_or_else(|| {
                    McpError::protocol(McpErrorCode::InvalidParams, "Missing parameters in tools/call request".to_string())
                })?;

                let result: McpClientInitializeResult = serde_json::from_value(params)?;
                let result = self.handler.initialize(result).await?;
                
                *self.initilized.write().await = true;

                Ok(McpResponse::success(request.id, Some(serde_json::to_value(result)?)))   
            }
            "shutdown" => {
                self.check_initialized(true).await?;
                self.handler.shutdown().await?;
                Ok(McpResponse::success(request.id, Some(serde_json::json!({}))))
            }
            "tools/list" => {
                self.check_initialized(true).await?;
                let tools_list = self.handler.tools_list().await?;
                Ok(McpResponse::success(request.id, Some(serde_json::json!({
                    "tools": tools_list
                }))))   
            }
            "tools/call" => {
                self.check_initialized(true).await?;
                match self.handler.tools_call(request.params.clone()).await {
                    Ok(result) =>
                        Ok(McpResponse::success(request.id, Some(serde_json::to_value(result)?))),
                    Err(err) => 
                        Ok(McpResponse::error(request.id, err.into()))
                }
            }
            _ => {
                self.check_initialized(true).await?;
                Ok(McpResponse::success(request.id, Some(serde_json::json!({}))))
            }
        }
    }

    async fn check_initialized(&self, state: bool) -> McpResult<()> {
        let initialized = *self.initilized.read().await;
        if state != initialized {
            return Err(McpError::protocol(
                McpErrorCode::ServerNotInitialized,
                "Server not initialized",
            ));
        }

        Ok(())
    }
}
